import numpy as np

import random
import time
import datetime
import RNG as RN_Gomoku
import Zobrist
start_time = time.time()
totaltime=0
#打印时间
def timeReset():
    global start_time
    start_time = time.time()
def printTime(tital=-1.0):
    global totaltime
    ss=time.time()-start_time
    totaltime+=float(ss);
    print("time used = " + str(datetime.timedelta(seconds=totaltime)))
    if(tital>ss):
        print("time remaining = " + str(datetime.timedelta(seconds=tital-ss)))
        return False
    else:
        return True
#打印棋盘
def printBoard(bd):
    print('   ',end='')
    for ii in range(15):
        print(' %d ' %(ii%10), end='')
    print('')
    for ii in range(15):
        print(' %d ' %(ii%10), end='')
        for jj in range(15):
            if(bd[ii,jj,0]==0 and bd[ii,jj,1]==0):
                if((ii,jj)==(7,7) or (ii==3 or ii==11) and (jj==3 or jj==11)):
                    print(' + ',end='')
                else:
                    print(' - ',end='')
            elif(bd[ii,jj,0]==1 and bd[ii,jj,1]==0):
                print('\033[96m O \033[0m',end='')
            elif(bd[ii,jj,0]==0 and bd[ii,jj,1]==1):
                print('\033[91m X \033[0m',end='')
            else:
                print(' E ',end='')
        print(' %d' %(ii%10))
    print('   ',end='')
    for ii in range(15):
        print(' %d ' %(ii%10), end='')
    print('')
def toTex(bd,ii,jj):
    if(bd[ii,jj,0]==0 and bd[ii,jj,1]==0):
        return ' - '
    elif(bd[ii,jj,0]==1 and bd[ii,jj,1]==0):
        return ' O '
    elif(bd[ii,jj,0]==0 and bd[ii,jj,1]==1):
        return ' X '
    else:
        return ' E '
#蒙特卡洛树节点
class Node(object):
    def __init__(self,prior: float,sidem: int):
        self.visit_count=0#访问次数
        self.to_play=sidem#哪一方先手
        self.prior=prior#先验概率
        self.value_sum=0.0#节点总评分
        self.actions=[]#可选择行为列表  0-225
        self.children=[]#子节点

    def set_state(self,prior: float,sidem: int):
        self.visit_count=0
        self.to_play=sidem
        self.prior=prior
        self.value_sum=0.0
        self.actions=[]
        self.children=[]

    def expanded(self):
        return len(self.children)>0

    def value(self):

        return self.value_sum/self.visit_count
board=np.array(np.zeros([15,15,2]), dtype=int)#棋盘
states = {}

movelist=[7*16 for i in range(225)]#记录移动列表，112 个即可判断游戏输赢
side2move=0 #哪一方先下棋
move_count=0 #下了多少步
num_simul=800#蒙特卡洛树模拟次数
#定义节点列表
nodes=[Node(0.1,1) for _ in range(max((num_simul+1)*224,500))]
nodes_used=0
def setN(nsimul):
    global num_simul,nodes
    num_simul=nsimul
    if((num_simul+1)*224>len(nodes)):
        nodes=[Node(0.1,1) for _ in range(max((num_simul+1)*224,500))]
    print("num_simul set to",num_simul)

#定义合法位置
def leagal_actions():
    # for i in range(15):
    #     for j in range(15):
    #         if(forbidden(i,j)==1 or forbidden(i,j)==3 or forbidden(i,j)==4):
    #             print("禁手!!", i, j)
    #             ban[i][j]=1;
    #         else:ban[i][j]=0;

    return [ii*15+jj for ii in range(15) for jj in range(15) if board[ii][jj][0]==0 and board[ii][jj][1]==0 ]
#返回0胜利或者X胜利 判断是否有胜利一方
def winLossDraw():
    sd=1-side2move
    posx,posy=movelist[move_count-1]//15,movelist[move_count-1]%15

    lx,lg=1,1
    while(posx-lx>=0 and board[posx-lx][posy][sd]==1):
        lg+=1
        lx+=1
    lx=1
    while(posx+lx<15 and board[posx+lx][posy][sd]==1):
        lg+=1
        lx+=1
    if(lg>=5):
        return side2move

    lx,lg=1,1
    while(posy-lx>=0 and board[posx][posy-lx][sd]==1):
        lg+=1
        lx+=1
    lx=1
    while(posy+lx<15 and board[posx][posy+lx][sd]==1):
        lg+=1
        lx+=1
    if(lg>=5):
        return side2move

    lx,lg=1,1
    while(posx-lx>=0 and posy-lx>=0 and board[posx-lx][posy-lx][sd]==1):
        lg+=1
        lx+=1
    lx=1
    while(posx+lx<15 and posy+lx<15 and board[posx+lx][posy+lx][sd]==1):
        lg+=1
        lx+=1
    if(lg>=5):
        return side2move

    lx,lg=1,1
    while(posx-lx>=0 and posy+lx<15 and board[posx-lx][posy+lx][sd]==1):
        lg+=1
        lx+=1
    lx=1
    while(posx+lx<15 and posy-lx>=0 and board[posx+lx][posy-lx][sd]==1):
        lg+=1
        lx+=1
    if(lg>=5):
        return side2move
    
    if(move_count==225):
        return 0.5
    return -1
#返回评估胜率
def evaluate(node: Node,evaluatePosition):
    global nodes_used
    wld=winLossDraw()

    if(wld!=-1):
        return wld
    policy_logits=evaluatePosition()
    node.actions=leagal_actions()
    policy=[policy_logits[i] for i in node.actions]
    policy_sum = sum(policy)
    num_p=len(node.actions)
    node.children=nodes[nodes_used:nodes_used+num_p]
    nodes_used+=num_p
    tp=1-node.to_play
    for i in range(num_p):
        node.children[i].set_state(policy[i]/policy_sum,tp)

    return policy_logits[-1]
#增加噪声
def add_exploration_noise(node: Node):
    noise=np.random.dirichlet(np.full(len(node.actions),.05))

    for i in range(len(node.actions)):
        node.children[i].prior = node.children[i].prior*.75+noise[i]*.25
fpuReduction=1.3
fpuReductionRoot=1.0
def setF(rdc,rdcr):
    global fpuReduction,fpuReductionRoot
    fpuReduction,fpuReductionRoot=rdc,rdcr
    print("FPU (normal/root) set to ",fpuReduction,fpuReductionRoot)
#UCB
def ucb_score(parent: Node, child: Node,isnotroot=True):
    pb_c=np.log((parent.visit_count+19653)/19652)+1.25
    pb_c*=np.sqrt(parent.visit_count)/(child.visit_count+1)
    prior_score=pb_c*child.prior
    if(child.visit_count==0):
        return (prior_score-fpuReduction+parent.value()) if isnotroot else (prior_score-fpuReductionRoot+parent.value())

    else:
        return prior_score-child.value()
#选择孩子节点 选择UCB最大的
def select_child(node: Node,isnotroot=True):
    _,action,child = max([(ucb_score(node,node.children[i],isnotroot),node.actions[i],node.children[i]) for i in range(len(node.actions))])

    return action,child
#更新Move后各个变量
def applyMove(pos):
    global board,side2move,move_count,movelist
    board[pos//15][pos%15][side2move]=1
    side2move=1-side2move
    movelist[move_count]=pos
    move_count+=1



    if(move_count%2==1):
     states[pos] =1;
    else:
     states[pos]=2;    #1代表为黑棋 ，0为白棋  pos:1or0
#反悔
def takeBack():
    global board,side2move,move_count
    move_count-=1
    pos=movelist[move_count]
    side2move=1-side2move
    board[pos//15][pos%15][side2move]=0
    states.pop(pos)

#反向传播 更新节点的value与count
def backpropagate(search_path, value: float):
    for node in search_path:
        node.value_sum += value if node.to_play == 0 else (1 - value)
        node.visit_count += 1
valueWt=0*num_simul
#得分
def actionScore(node:Node):
    if(node.visit_count<1):
        return -valueWt
    else:
        return node.visit_count-valueWt*node.value()
def move_to_location(move):
        '''
        transfer move number to coordinate

        3*3 board's moves like:
        6 7 8
        3 4 5
        0 1 2
        and move 5's location is (1,2)
        '''
        h = move //15
        w = move %15
        return [h, w]
#选择哪一步30步以后增加一定噪声，探索性
def select_action(root: Node,add_noise=True):
    if(move_count>30 or (not add_noise)):

         visit_counts = [(actionScore(root.children[i]),root.actions[i]) for i in range(len(root.actions))]
         #print(visit_counts) #0-225可选位置的评分 （-50，0） （-50，1） 0，1位置评分为-50
         _, action = max(visit_counts)
         print(action)
         v=visit_counts.copy()
         v.sort()
         print(v)
         c=1
         _,move=v[-c]
         i, j = move_to_location(action)
         ban = forbidden(i, j)

         while ban == 1 or ban == 3 or ban == 4:
             print("禁手!!", i, j,ban)
             c=c+1
             a,move =v[-c]
             if(a==-50):
                _,move=v[-c+1]
                break;

             i, j = move_to_location(move)
             ban = forbidden(i, j)
         action=move
         '''
          i, j = move_to_location(action)
          while ban == 1 or ban == 3 or ban == 4:
                    print("禁手!!", i, j)
                    self.board.availables.remove(move)
                    move, move_probs = player_in_turn.get_action(self.board, is_selfplay=False, print_probs_value=1)
                    i, j = board.move_to_location(move)
                    ban = forbidden(i, j)'''
    else:
        visit_counts = np.array([(root.children[i].visit_count)**1.5 for i in range(len(root.actions))])
        visit_counts_sum=np.sum(visit_counts)
        action = np.random.choice(root.actions,p=visit_counts/visit_counts_sum)


    pol_map=np.zeros([15,15])
    visit_counts = np.array([root.children[i].visit_count for i in range(len(root.actions))])
    visit_counts_sum=np.sum(visit_counts)
    for i in range(len(root.actions)):
        pol_map[root.actions[i]//15][root.actions[i]%15]=visit_counts[i]/visit_counts_sum
    return action,pol_map,root.value()
#运行
def run_mcts(evaluatePosition,add_noise=True):
    global nodes_used
    nodes_used=0
    root = Node(1.0,side2move)
    root.visit_count=1
    root.value_sum=evaluate(root,evaluatePosition) if side2move==0 else 1.0-evaluate(root,evaluatePosition)
    depth=0
    maxdepth=0
    if(add_noise):
        add_exploration_noise(root)
#当前节点为叶子节点
    for _ in range(num_simul):
        node=root
        search_path=[node]
        depth=0
#扩展节点
        while node.expanded():
            depth+=1

            action,node=select_child(node,depth>1)
            applyMove(action)
            search_path.append(node)

        value=evaluate(node,evaluatePosition)
#反向传播更新
        backpropagate(search_path,value)
        for __ in range(depth):
            takeBack()

        maxdepth=max(maxdepth,depth)
#返回选择棋盘位置与胜率
    return select_action(root,add_noise),maxdepth
#加载引擎
def loadEngine(idx,wtf,hszs=500000):
    global a0eng,a0eng2,hashTB,hashTB2
    if(idx==1):
        a0eng =RN_Gomoku.A0_ENG(64,wtf)
        hashTB=Zobrist.ZobristHash(hszs)
    elif(idx==2):
        a0eng2=RN_Gomoku.A0_ENG(64,wtf)
        hashTB2=Zobrist.ZobristHash(hszs)
    else:
        print("Engine index must be 1 or 2!")

#神经网络预测
def evaluatePositionA():# nn_old
    kkey,rstt=hashTB.getValue(board)
    if(len(rstt)>0):
        return rstt

    else:
        if(side2move==1):

            rstt=a0eng.a0_eng(np.array([np.flip(board,-1)],dtype='<f4'),training=False)[0].numpy()
            rstt[-1]=1-rstt[-1]
        else:
            rstt=a0eng.a0_eng(np.array([board.copy()],dtype='<f4'),training=False)[0].numpy()
        hashTB.setValue(kkey,board,rstt)
        return rstt




def forbidden(h, w,player=1):
        height =15
        width = 15
        n = 5
        restricted = [0,0] # 33, 44
        buffer = ["33333333333" for j in range(4)] # 3: 棋盤外
        print(states)
        for i in range(w - n, w + n + 1): #下至上
            if i < 0 or i >= width:
                continue
            buffer[0] = buffer[0][:i-w+5] + str(states.get(i + h * width, 0)) + buffer[0][i-w+6:]
            buffer[0] = buffer[0][:5] + str(player) + buffer[0][6:]

        for i in range(h - n, h + n + 1): #左至右
            if i < 0 or i >= height:
                continue
            buffer[1] = buffer[1][:i-h+5] + str(states.get(i * width + w, 0)) + buffer[1][i-h+6:]
            buffer[1] = buffer[1][:5] + str(player) + buffer[1][6:]

        i = h - n - 1
        j = w + n + 1
        for count in range(11): #左上至右下
            i += 1
            j -= 1
            if i < 0 or i >= height or j < 0 or j >= width:
                continue
            buffer[2] = buffer[2][:count] + str(states.get(i * width + j, 0)) + buffer[2][count+1:]
            buffer[2] = buffer[2][:5] + str(player) + buffer[2][6:]

        i = h - n - 1
        j = w - n - 1
        for count in range(11): #左下至右上
            i += 1
            j += 1
            if i < 0 or i >= height or j < 0 or j >= width:
                continue
            buffer[3] = buffer[3][:count] + str(states.get(i * width + j, 0)) + buffer[3][count+1:]
            buffer[3] = buffer[3][:5] + str(player) + buffer[3][6:]

        """print()
        print(buffer[0])
        print(buffer[1])
        print(buffer[2])
        print(buffer[3])
        print()"""

        for i in range(4):
            judge = 0
            for j in range(11):
                judge = judge+1 if(buffer[i][j]==str(player)) else 0
                if judge == 5:
                    return 1 if (j<10 and buffer[i][j+1]==str(player)) else 2 # 1長連禁手  2 win

        for i in range(4):
            before = restricted[1]
            restricted[1] = restricted[1]+1 if(buffer[i].find("011110")!=-1 or buffer[i].find("011112")!=-1 or buffer[i].find("211110")!=-1 or buffer[i].find("011113")!=-1 or buffer[i].find("311110")!=-1) else restricted[1]
            #if buffer[i].find("10111") != -1:
            #    restricted[1] = restricted[1] if(buffer[i].find("1101110")!=-1 or buffer[i].find("0101111")!=-1 or buffer[i].find("1101111")!=-1) else restricted[1]+1
            restricted[1] = restricted[1]+1 if(buffer[i].find("10111") != -1) else restricted[1]
            if buffer[i].find("11011") != -1:
                if(not (buffer[i].find("1110110")!=-1 or buffer[i].find("0110111")!=-1 or buffer[i].find("1110111")!=-1)):
                    restricted[1] = restricted[1]+1
                    #print(i, buffer[i].find("11011"))
                    if(not (buffer[i].find("1110110", buffer[i].find("11011")+1)!=-1 or buffer[i].find("0110111", buffer[i].find("11011")+1)!=-1 or buffer[i].find("1110111", buffer[i].find("11011")+1)!=-1)):
                        if buffer[i].find("11011", buffer[i].find("11011")+1) != -1:
                            restricted[1] = restricted[1]+1
                            #print(i, buffer[i].find("11011", buffer[i].find("11011")+1))
            #if buffer[i].find("11101") != -1:
            #    restricted[1] = restricted[1] if(buffer[i].find("1111010")!=-1 or buffer[i].find("0111011")!=-1 or buffer[i].find("1111011")!=-1) else restricted[1]+1
            restricted[1] = restricted[1]+1 if(buffer[i].find("11101") != -1) else restricted[1]
            """print('before: ', restricted[1])"""
            if restricted[1] - before > 1:
                restricted[1] = restricted[1] if (buffer[i].find("111010111")!=-1 or buffer[i].find("11011011")!=-1 or buffer[i].find("1011101")!=-1) else 1
                if restricted[1] > 1:
                    return 3
            """print('after: ', restricted[1])"""

        #print("四:{}".format(restricted[1]))
        if restricted[1] > 1:
            return 3

        for i in range(4):
            if buffer[i].find("01110") != -1:
                restricted[0] = restricted[0] if(buffer[i].find("100111001")!=-1 or buffer[i].find("11101")!=-1 or buffer[i].find("10111")!=-1 or buffer[i].find("10011102")!=-1 or buffer[i].find("20111001")!=-1 or buffer[i].find("2011102")!=-1 or buffer[i].find("2011103")!=-1 or buffer[i].find("3011102")!=-1 or buffer[i].find("10011103")!=-1 or buffer[i].find("30111001")!=-1) else restricted[0]+1
            if buffer[i].find("010110") != -1:
                restricted[0] = restricted[0] if(buffer[i].find("00101101")!=-1 or buffer[i].find("10101100")!=-1 or buffer[i].find("10101101")!=-1 or buffer[i].find("10101102")!=-1 or buffer[i].find("20101101")!=-1 or buffer[i].find("10101103")!=-1 or buffer[i].find("30101101")!=-1) else restricted[0]+1
            if buffer[i].find("011010") != -1:
                restricted[0] = restricted[0] if(buffer[i].find("00110101")!=-1 or buffer[i].find("10110100")!=-1 or buffer[i].find("10110101")!=-1 or buffer[i].find("10110102")!=-1 or buffer[i].find("20110101")!=-1 or buffer[i].find("10110103")!=-1 or buffer[i].find("30110101")!=-1) else restricted[0]+1

        #print("三:{}".format(restricted[0]))
        #print()
        if(restricted[0] > 1):
            return 4

